import tensorflow as tf
from tensorflow import keras
import numpy as np
# import sklearn


class TestModel(keras.Model):
    def __init__(self, num_classes=10):
        super(TestModel, self).__init__(name="test_model")
        self.num_classes = num_classes
        self.dense1 = keras.layers.Dense(32, activation='relu')
        self.dense2 = keras.layers.Dense(num_classes, activation='sigmoid')

    def call(self, inputs, training=None, mask=None):
        x = self.dense1(inputs)
        x = self.dense2(x)
        return x

    def compute_output_shape(self, input_shape):
        shape = tf.TensorShape(input_shape).as_list()
        shape[-1] = self.num_classes
        return tf.TensorShape(shape)


if __name__ == '__main__':
    model = TestModel(10)
    model.compile(optimizer=keras.optimizers.RMSprop(0.001),
                  loss="categorical_crossentropy",
                  metrics=['accuracy'])

    data = np.random.random((1000, 32))
    labels = np.random.random((1000, 10))

    model.fit(data, labels, batch_size=32, epochs=5)

    model.evaluate(data, labels, batch_size=32)

